#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Author  : 邢建辉
# @Email   : xjh_0125@sina.com
# @Time    : 2022/5/12 14:30
# @Software: PyCharm
# @File    : l703KthLargest.py


class KthLargest:
    def __init__(self, k: int, nums: list[int]):
        self.arr = []
        self.k = k
        self.arr = nums[:k]
        self.build_heap()
        for i in nums[k:]:
            self.add(i)

    def swap(self, l, r):
        self.arr[l], self.arr[r] = self.arr[r], self.arr[l]

    def build_heap(self):
        # mid = len(self.arr) // 2
        # for i in range(mid, -1, -1):
        #     self.heapify(i)
        if len(self.arr) < 2:
            return
        for i in range(len(self.arr)):
            self.heapify_up(i)

    def heapify_up(self, pos):
        while pos > 0:
            parent = (pos - 1) // 2
            if self.arr[pos] < self.arr[parent]:
                self.swap(pos, parent)
                pos = parent
            else:
                break

    def heapify(self, pos):
        '''
        构建最小堆 跟节点最小，左子树大于右子树
        按照索引排
              0
           1  *  2
         3  4 * 5  6
        :param arr:生成堆的源数组
        :param heap_size: 堆大小 小于等于len(arr)
        :param pos:堆顶点位置
        :return:
        '''
        l, r = pos * 2 + 1, pos * 2 + 2
        k = len(self.arr)
        while l < k:
            largest_child = l
            if r < k and self.arr[l] > self.arr[r]:
                largest_child = r
            if self.arr[pos] <= self.arr[largest_child]:
                break
            else:
                self.swap(pos, largest_child)
                pos = largest_child  # 此时largest为左或右子树，再遍历当前节点的左右子树
                l, r = pos * 2 + 1, pos * 2 + 2

    def add(self, val: int) -> int:
        if len(self.arr) < self.k:
            self.arr.append(val)
            self.build_heap()
        elif val > self.arr[0]:
            self.arr[0] = val
            self.build_heap()
        return self.arr[0]


import heapq


class KthLargest2:
    def __init__(self, k: int, nums: list[int]):
        # self.k = k
        # self.size = min(len(nums), k)
        # self.q = nums[0:k]
        # heapq.heapify(self.q)
        # for i in nums[k:]:
        #     self.add(i)
        self.heap = nums
        self.k = k
        heapq.heapify(self.heap)
        while len(self.heap) > k:
            heapq.heappop(self.heap)

    def add(self, val: int) -> int:
        if len(self.heap) < self.k:
            heapq.heappush(self.heap, val)
        elif val > self.heap[0]:
            heapq.heapreplace(self.heap, val)
        return self.heap[0]


if __name__ == '__main__':
    arr = [0]
    k = 2
    s = KthLargest2(k, arr)
    print(s.add(-1))
    print(s.add(1))
    print(s.add(-2))
    print(s.add(-4))
    print(s.add(3))
